package com.ztom.top100;

import java.util.HashMap;
import java.util.Map;

/**
 * 路径总和 3
 * <p>
 * https://leetcode-cn.com/problems/path-sum-iii/
 *
 * @author ZhangTao
 */
public class Code88PathSum {

    public class TreeNode {
        int val;
        TreeNode left;
        TreeNode right;

        TreeNode() {
        }

        TreeNode(int val) {
            this.val = val;
        }

        TreeNode(int val, TreeNode left, TreeNode right) {
            this.val = val;
            this.left = left;
            this.right = right;
        }
    }


    public int pathSum1(TreeNode root, int targetSum) {
        if (root == null) {
            return 0;
        }

        // 以每个节点为起始节点, 向下寻找能否凑齐 target
        int res = sum(root, targetSum);
        // 遍历左树
        res += pathSum1(root.left, targetSum);
        // 遍历右树
        res += pathSum1(root.right, targetSum);
        return res;
    }

    private int sum(TreeNode root, int rest) {
        if (root == null) {
            return 0;
        }

        int res = 0;
        if (root.val == rest) {
            // 剩余为 0 说明找到一种方法
            res++;
        }
        // 向左树走
        res += sum(root.left, rest - root.val);
        // 向右树走
        res += sum(root.right, rest - root.val);
        return res;
    }

    /**
     * 前缀和 + 回溯
     *
     * @param root
     * @param targetSum
     * @return
     */
    public int pathSum(TreeNode root, int targetSum) {
        if (root == null) {
            return 0;
        }

        // 记录前缀和出现的次数
        Map<Integer, Integer> prefixMap = new HashMap<>();
        prefixMap.put(0, 1);
        return dfs(root, prefixMap, 0, targetSum);
    }

    private int dfs(TreeNode node, Map<Integer, Integer> prefixMap, int curSum, int targetSum) {
        if (node == null) {
            return 0;
        }

        int res = 0;
        // 累加当前节点值
        curSum += node.val;

        // 前缀和, curNum - x = targetSum 即需要看 x 是否记录过
        // 可以从任意节点出发, 所以记录出现的次数, 就是符合条件的数量
        res = prefixMap.getOrDefault(curSum - targetSum, 0);
        // 记录当前前缀和
        prefixMap.put(curSum, prefixMap.getOrDefault(curSum, 0) + 1);
        // 去左右子树查找
        res += dfs(node.left, prefixMap, curSum, targetSum);
        res += dfs(node.right, prefixMap, curSum, targetSum);
        // 回溯
        prefixMap.put(curSum, prefixMap.getOrDefault(curSum, 0) - 1);

        return res;
    }
}
